from random import Random

import setting
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


# AES ECB mode without IV

def encrypt(key: str, data: str) -> str:
    key = key.encode('utf-8')
    data = pad(data.encode(), AES.block_size)
    cipher = AES.new(key, AES.MODE_ECB)
    rst = base64.b64encode(cipher.encrypt(data))
    return rst.decode("utf-8", 'ignore')


def decrypt(key: str, enc: str) -> str:
    enc = enc.encode("utf-8")
    enc = base64.b64decode(enc)
    cipher = AES.new(key.encode('utf-8'), AES.MODE_ECB)
    return unpad(cipher.decrypt(enc), AES.block_size).decode()


def generateAESKey(param_custom_key: str) -> str:
    RES = ''
    if len(param_custom_key) < 16:
        RES += create_salt(16 - len(param_custom_key))
    else:
        RES = param_custom_key[0:16]
    return RES


def create_salt(length=4):
    salt = ''
    chars = 'AaBbCcDdEeFfGgHhIiJjKkLlMmNnOoPpQqRrSsTtUuVvWwXxYyZz0123456789+-*'
    len_chars = len(chars) - 1
    random = Random()
    for i in range(length):
        # 每次从chars中随机取一位
        salt += chars[random.randint(0, len_chars)]
    return salt


if __name__ == '__main__':
    encrypted = encrypt(key=setting.security.app_default_AES_key, data="U2FsdGVkX1+rXVAKnwaaxp0h7AKtOZhiJTv9jEf9R4I="
                        )
    # print('[myAES.py][main]encrypted ECB Base64:', encrypted)
    decrypted = decrypt(key="hello_next_hello",
                        enc="U2FsdGVkX1/Jf2N8ksOhqH4Jfk4yz52enJMu1qqsFVo="

                        )

    print('[myAES.py][main]decrypted data:', decrypted)
